use axum::{routing::post, Router};
use axum::extract::DefaultBodyLimit;
use std::net::SocketAddr;
use tokio::sync::RwLock;
use vdb::{VectorDB, Vector};

async fn run_server(db: std::sync::Arc<RwLock<VectorDB>>, ip: String, port: u16) {
    let app = Router::new()
        .route("/create_collection", post(handle_create_collection))
        .route("/insert", post(handle_insert))
        .route("/delete", post(handle_delete))
        .route("/search", post(handle_search))
        .route("/close", post(handle_close_collection))
        .layer(DefaultBodyLimit::max(10 * 1024 * 1024))
        .with_state(db);

    let addr = match ip.parse() {
        Ok(parsed_ip) => SocketAddr::new(parsed_ip, port),
        Err(err) => {
            log::error!("Invalid IP address '{}': {}", ip, err);
            return;
        }
    };

    let listener = match tokio::net::TcpListener::bind(&addr).await {
        Ok(listener) => listener,
        Err(err) => {
            log::error!("Failed to bind to address {}: {}", addr, err);
            return;
        }
    };

    log::info!("Server running on {}", addr);

    if let Err(err) = axum::serve(listener, app).await {
        log::error!("Server error: {}", err);
    }
}

#[derive(clap::Parser)]
struct Args {
    /// Storage format (proto|json)
    #[clap(long, default_value = "proto")]
    storage_type: String,

    /// IP address to bind the server
    #[clap(long, default_value = "127.0.0.1")]
    ip: String,

    /// Port to bind the server
    #[clap(long, default_value = "3000")]
    port: u16,
}

#[tokio::main]
async fn main() -> std::io::Result<()> {
    std::fs::create_dir_all("logs").expect("创建日志目录失败");
    std::fs::create_dir_all("collections").expect("创建集合目录失败");

    fern::Dispatch::new()
        .format(|out, message, record| {
            out.finish(format_args!(
                "{} [{}] {} - {}:{}",
                chrono::Local::now().format("%Y-%m-%d %H:%M:%S"),
                record.level(),
                message,
                record.file().unwrap_or("unknown"),
                record.line().unwrap_or(0)
            ))
        })
        .chain(
            fern::Dispatch::new()
                .chain(std::io::stdout())
                .level(log::LevelFilter::Info)
        )
        .chain(
            fern::Dispatch::new()
                .chain(fern::log_file(format!("logs/{}.log", chrono::Local::now().format("%Y-%m-%d"))).expect("创建日志文件失败"))
                .level(log::LevelFilter::Info)
        )
        .apply()
        .expect("初始化日志系统失败");
    log::info!("Starting vector database server");

    let args: Args = clap::Parser::parse();
    let db = VectorDB::new(args.storage_type.clone());
    db.initialize("./collections").await?; // Ensure collections are loaded from disk
    let db = std::sync::Arc::new(RwLock::new(db));

    run_server(db, args.ip, args.port).await;

    Ok(())
}

async fn handle_insert(
    axum::extract::State(db): axum::extract::State<std::sync::Arc<tokio::sync::RwLock<VectorDB>>>,
    axum::Json(payload): axum::Json<InsertRequest>,
) -> axum::http::StatusCode {
    let InsertRequest { collection_name, vectors } = payload;
    log::info!(
        "Received insert request for collection '{}' with {} vectors",
        collection_name,
        vectors.len()
    );

    let vector_objects: Vec<Vector> = vectors
        .into_iter()
        .map(|vector| Vector {
            id: vector.id,
            embedding: vector.embedding,
            metadata: vector.metadata,
        })
        .collect();

    if let Err(err) = db
        .write()
        .await
        .bulk_insert(collection_name.clone(), vector_objects)
        .await
    {
        log::error!("Failed to bulk insert vectors: {}", err);
        return axum::http::StatusCode::INTERNAL_SERVER_ERROR;
    }

    axum::http::StatusCode::OK
}

#[derive(serde::Deserialize)]
struct InsertRequest {
    collection_name: String,
    vectors: Vec<VectorPayload>,
}

#[derive(serde::Deserialize)]
struct VectorPayload {
    id: String,
    embedding: Vec<f32>,
    metadata: serde_json::Value,
}

async fn handle_delete(
    axum::extract::State(db): axum::extract::State<std::sync::Arc<tokio::sync::RwLock<VectorDB>>>,
    axum::Json(payload): axum::Json<DeleteRequest>,
) -> axum::http::StatusCode {
    let DeleteRequest { collection_name, ids } = payload;
    log::info!(
        "Received delete request for collection '{}' with {} IDs",
        collection_name,
        ids.len()
    );

    if let Err(err) = db.write().await.bulk_delete(collection_name.clone(), ids).await {
        log::error!("Failed to bulk delete vectors from collection '{}': {}", collection_name, err);
        return axum::http::StatusCode::INTERNAL_SERVER_ERROR;
    }

    axum::http::StatusCode::OK
}

#[derive(serde::Deserialize)]
struct DeleteRequest {
    collection_name: String,
    ids: Vec<String>,
}

async fn handle_search(
    axum::extract::State(db): axum::extract::State<std::sync::Arc<tokio::sync::RwLock<VectorDB>>>,
    axum::Json(payload): axum::Json<SearchRequest>,
) -> axum::Json<Result<Vec<(String, f32)>, String>> {
    let SearchRequest { collection_name, query, k } = payload;
    log::info!(
        "Received search request for collection '{}' with vector length: {}",
        collection_name,
        query.len()
    );

    let results = db.read().await.search(collection_name, query, k).await;
    axum::Json(results)
}

#[derive(serde::Deserialize)]
struct SearchRequest {
    collection_name: String,
    query: Vec<f32>,
    k: usize,
}

async fn handle_create_collection(
    axum::extract::State(db): axum::extract::State<std::sync::Arc<tokio::sync::RwLock<VectorDB>>>,
    axum::Json(payload): axum::Json<CreateCollectionRequest>,
) -> axum::http::StatusCode {
    log::info!("Received request to create collection '{}'", payload.name);

    db.write().await.create_collection(&payload.name).await;

    axum::http::StatusCode::OK
}

async fn handle_close_collection(
    axum::extract::State(db): axum::extract::State<std::sync::Arc<tokio::sync::RwLock<VectorDB>>>,
    axum::Json(payload): axum::Json<CloseCollectionRequest>,
) -> axum::http::StatusCode {
    log::info!("Received request to close collection '{}'", payload.name);

    match db.write().await.close_collection(&payload.name).await {
        Ok(_) => axum::http::StatusCode::OK,
        Err(err) => {
            log::error!("Failed to close collection '{}': {}", payload.name, err);
            axum::http::StatusCode::INTERNAL_SERVER_ERROR
        }
    }
}

#[derive(serde::Deserialize)]
struct CloseCollectionRequest {
    name: String,
}

#[derive(serde::Deserialize)]
struct CreateCollectionRequest {
    name: String,
}

